<!DOCTYPE html>
<html lang="zh-cn" dir="ltr">
    <head><meta charset='utf-8'>
<meta name='viewport' content='width=device-width, initial-scale=1'><meta name='description' content='加载一个很大的数据集，全部加载到内存的话会导致内存爆炸（OOM），该怎么办？'><title>pytorch加载大数据集</title>

<link rel='canonical' href='https://charent.github.io/p/pytorch%E5%8A%A0%E8%BD%BD%E5%A4%A7%E6%95%B0%E6%8D%AE%E9%9B%86/'>

<link rel="stylesheet" href="/scss/style.min.65463afd28d606277b44441c8bbb8b0277823a2d0c03ab8ba9d0567d274f7b43.css"><meta property='og:title' content='pytorch加载大数据集'>
<meta property='og:description' content='加载一个很大的数据集，全部加载到内存的话会导致内存爆炸（OOM），该怎么办？'>
<meta property='og:url' content='https://charent.github.io/p/pytorch%E5%8A%A0%E8%BD%BD%E5%A4%A7%E6%95%B0%E6%8D%AE%E9%9B%86/'>
<meta property='og:site_name' content='Charent的博客'>
<meta property='og:type' content='article'><meta property='article:section' content='Post' /><meta property='article:published_time' content='2023-11-05T00:00:00&#43;00:00'/><meta property='article:modified_time' content='2023-11-05T00:00:00&#43;00:00'/>
<meta name="twitter:title" content="pytorch加载大数据集">
<meta name="twitter:description" content="加载一个很大的数据集，全部加载到内存的话会导致内存爆炸（OOM），该怎么办？">
    </head>
    <body class="
    article-page
    ">
    <script>
        (function() {
            const colorSchemeKey = 'StackColorScheme';
            if(!localStorage.getItem(colorSchemeKey)){
                localStorage.setItem(colorSchemeKey, "auto");
            }
        })();
    </script><script>
    (function() {
        const colorSchemeKey = 'StackColorScheme';
        const colorSchemeItem = localStorage.getItem(colorSchemeKey);
        const supportDarkMode = window.matchMedia('(prefers-color-scheme: dark)').matches === true;

        if (colorSchemeItem == 'dark' || colorSchemeItem === 'auto' && supportDarkMode) {
            

            document.documentElement.dataset.scheme = 'dark';
        } else {
            document.documentElement.dataset.scheme = 'light';
        }
    })();
</script>
<div class="container main-container flex on-phone--column extended"><aside class="sidebar left-sidebar sticky ">
    <button class="hamburger hamburger--spin" type="button" id="toggle-menu" aria-label="切换菜单">
        <span class="hamburger-box">
            <span class="hamburger-inner"></span>
        </span>
    </button>

    <header>
        
            
            <figure class="site-avatar">
                <a href="/">
                
                    
                    
                    
                        
                        <img src="/img/avatar_hu544a9780599d84daff7ca906f17dd84c_13353_300x0_resize_box_3.png" width="300"
                            height="300" class="site-logo" loading="lazy" alt="Avatar">
                    
                
                </a>
                
                    <span class="emoji">😆</span>
                
            </figure>
            
        
        
        <div class="site-meta">
            <h1 class="site-name"><a href="/">Charent的博客</a></h1>
            <h2 class="site-description">不积硅步，无以至千里</h2>
        </div>
    </header><ol class="social-menu">
            
                <li>
                    <a 
                        href='https://github.com/charent'
                        target="_blank"
                        title="GitHub"
                    >
                        
                        
                            <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-brand-github" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
  <path d="M9 19c-4.3 1.4 -4.3 -2.5 -6 -3m12 5v-3.5c0 -1 .1 -1.4 -.5 -2c2.8 -.3 5.5 -1.4 5.5 -6a4.6 4.6 0 0 0 -1.3 -3.2a4.2 4.2 0 0 0 -.1 -3.2s-1.1 -.3 -3.5 1.3a12.3 12.3 0 0 0 -6.2 0c-2.4 -1.6 -3.5 -1.3 -3.5 -1.3a4.2 4.2 0 0 0 -.1 3.2a4.6 4.6 0 0 0 -1.3 3.2c0 4.6 2.7 5.7 5.5 6c-.6 .6 -.6 1.2 -.5 2v3.5" />
</svg>



                        
                    </a>
                </li>
            
        </ol><ol class="menu" id="main-menu">
        
        
        

        <li >
            <a href='/' >
                
                
                
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-home" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <polyline points="5 12 3 12 12 3 21 12 19 12" />
  <path d="M5 12v7a2 2 0 0 0 2 2h10a2 2 0 0 0 2 -2v-7" />
  <path d="M9 21v-6a2 2 0 0 1 2 -2h2a2 2 0 0 1 2 2v6" />
</svg>



                
                <span>主页</span>
            </a>
        </li>
        
        

        <li >
            <a href='/archives/' >
                
                
                
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-archive" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <rect x="3" y="4" width="18" height="4" rx="2" />
  <path d="M5 8v10a2 2 0 0 0 2 2h10a2 2 0 0 0 2 -2v-10" />
  <line x1="10" y1="12" x2="14" y2="12" />
</svg>



                
                <span>归档</span>
            </a>
        </li>
        
        

        <li >
            <a href='/search/' >
                
                
                
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-search" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <circle cx="10" cy="10" r="7" />
  <line x1="21" y1="21" x2="15" y2="15" />
</svg>



                
                <span>搜索</span>
            </a>
        </li>
        
        

        <li >
            <a href='/%E5%85%B3%E4%BA%8E/' >
                
                
                
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-link" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <path d="M10 14a3.5 3.5 0 0 0 5 0l4 -4a3.5 3.5 0 0 0 -5 -5l-.5 .5" />
  <path d="M14 10a3.5 3.5 0 0 0 -5 0l-4 4a3.5 3.5 0 0 0 5 5l.5 -.5" />
</svg>



                
                <span>关于</span>
            </a>
        </li>
        

        <div class="menu-bottom-section">
                <li id="i18n-switch">  
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-language" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
  <path d="M4 5h7" />
  <path d="M9 3v2c0 4.418 -2.239 8 -5 8" />
  <path d="M5 9c-.003 2.144 2.952 3.908 6.7 4" />
  <path d="M12 20l4 -9l4 9" />
  <path d="M19.1 18h-6.2" />
</svg>



                    <select name="language" onchange="window.location.href = this.selectedOptions[0].value">
                        
                            <option value="https://charent.github.io/" selected>中文</option>
                        
                            <option value="https://charent.github.io/en/" >English</option>
                        
                    </select>
                </li>
            
            
            
                <li id="dark-mode-toggle">
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-toggle-left" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <circle cx="8" cy="12" r="2" />
  <rect x="2" y="6" width="20" height="12" rx="6" />
</svg>



                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-toggle-right" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <circle cx="16" cy="12" r="2" />
  <rect x="2" y="6" width="20" height="12" rx="6" />
</svg>



                    <span>暗色模式</span>
                </li>
            
        </div>
    </ol>
</aside>
<main class="main full-width">
    <article class="main-article">
    <header class="article-header">

    <div class="article-details">
    
    <header class="article-category">
        
            <a href="/categories/%E6%B7%B1%E5%BA%A6%E5%AD%A6%E4%B9%A0/" >
                深度学习
            </a>
        
            <a href="/categories/pytorch/" >
                pytorch
            </a>
        
            <a href="/categories/%E6%95%B0%E6%8D%AE%E9%9B%86/" >
                数据集
            </a>
        
    </header>
    

    <div class="article-title-wrapper">
        <h2 class="article-title">
            <a href="/p/pytorch%E5%8A%A0%E8%BD%BD%E5%A4%A7%E6%95%B0%E6%8D%AE%E9%9B%86/">pytorch加载大数据集</a>
        </h2>
    
        
        <h3 class="article-subtitle">
            加载一个很大的数据集，全部加载到内存的话会导致内存爆炸（OOM），该怎么办？
        </h3>
        
    </div>

    
    <footer class="article-time">
        
            <div>
                <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-calendar-time" width="56" height="56" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <path d="M11.795 21h-6.795a2 2 0 0 1 -2 -2v-12a2 2 0 0 1 2 -2h12a2 2 0 0 1 2 2v4" />
  <circle cx="18" cy="18" r="4" />
  <path d="M15 3v4" />
  <path d="M7 3v4" />
  <path d="M3 11h16" />
  <path d="M18 16.496v1.504l1 1" />
</svg>
                <time class="article-time--published">2023-11-05</time>
            </div>
        

        
    </footer>
    

    
</div>
</header>

    <section class="article-content">
    
    
    <h1 id="问题背景">问题背景</h1>
<p>假设我们有一个很大的数据集（GB级别），需要加载到内存中，但是全部加载的话会导致内存爆炸。同时我们还要保证可以在训练的不同的epoch打乱数据，所以先打乱数据保存到文件，再一行一行（一个样本一个样本）加载进来行不通，而且一行一行加载会导致非常大的IO，速度还会变慢。</p>
<h1 id="解决办法">解决办法</h1>
<ol>
<li>先把数据集转换为<code>parquet</code>文件，这中个格式的好处是可以分块读取，有效降低磁盘IO。</li>
<li>在<code>torch</code>数据集类<code>dataset</code>中对<code>parquet</code>文件循环，一次加载<code>N</code>条数据到内存缓冲区中，对缓存中的N条数据执行打乱（shuffle）操作即可</li>
<li>为了实现每个<code>epoch</code>都可以循环区数据集中的数据，可以使用<code>python</code>的<code>yield</code>特性实现迭代。</li>
</ol>
<h1 id="代码示例">代码示例</h1>
<div class="highlight"><div class="chroma">
<table class="lntable"><tr><td class="lntd">
<pre tabindex="0" class="chroma"><code><span class="lnt">  1
</span><span class="lnt">  2
</span><span class="lnt">  3
</span><span class="lnt">  4
</span><span class="lnt">  5
</span><span class="lnt">  6
</span><span class="lnt">  7
</span><span class="lnt">  8
</span><span class="lnt">  9
</span><span class="lnt"> 10
</span><span class="lnt"> 11
</span><span class="lnt"> 12
</span><span class="lnt"> 13
</span><span class="lnt"> 14
</span><span class="lnt"> 15
</span><span class="lnt"> 16
</span><span class="lnt"> 17
</span><span class="lnt"> 18
</span><span class="lnt"> 19
</span><span class="lnt"> 20
</span><span class="lnt"> 21
</span><span class="lnt"> 22
</span><span class="lnt"> 23
</span><span class="lnt"> 24
</span><span class="lnt"> 25
</span><span class="lnt"> 26
</span><span class="lnt"> 27
</span><span class="lnt"> 28
</span><span class="lnt"> 29
</span><span class="lnt"> 30
</span><span class="lnt"> 31
</span><span class="lnt"> 32
</span><span class="lnt"> 33
</span><span class="lnt"> 34
</span><span class="lnt"> 35
</span><span class="lnt"> 36
</span><span class="lnt"> 37
</span><span class="lnt"> 38
</span><span class="lnt"> 39
</span><span class="lnt"> 40
</span><span class="lnt"> 41
</span><span class="lnt"> 42
</span><span class="lnt"> 43
</span><span class="lnt"> 44
</span><span class="lnt"> 45
</span><span class="lnt"> 46
</span><span class="lnt"> 47
</span><span class="lnt"> 48
</span><span class="lnt"> 49
</span><span class="lnt"> 50
</span><span class="lnt"> 51
</span><span class="lnt"> 52
</span><span class="lnt"> 53
</span><span class="lnt"> 54
</span><span class="lnt"> 55
</span><span class="lnt"> 56
</span><span class="lnt"> 57
</span><span class="lnt"> 58
</span><span class="lnt"> 59
</span><span class="lnt"> 60
</span><span class="lnt"> 61
</span><span class="lnt"> 62
</span><span class="lnt"> 63
</span><span class="lnt"> 64
</span><span class="lnt"> 65
</span><span class="lnt"> 66
</span><span class="lnt"> 67
</span><span class="lnt"> 68
</span><span class="lnt"> 69
</span><span class="lnt"> 70
</span><span class="lnt"> 71
</span><span class="lnt"> 72
</span><span class="lnt"> 73
</span><span class="lnt"> 74
</span><span class="lnt"> 75
</span><span class="lnt"> 76
</span><span class="lnt"> 77
</span><span class="lnt"> 78
</span><span class="lnt"> 79
</span><span class="lnt"> 80
</span><span class="lnt"> 81
</span><span class="lnt"> 82
</span><span class="lnt"> 83
</span><span class="lnt"> 84
</span><span class="lnt"> 85
</span><span class="lnt"> 86
</span><span class="lnt"> 87
</span><span class="lnt"> 88
</span><span class="lnt"> 89
</span><span class="lnt"> 90
</span><span class="lnt"> 91
</span><span class="lnt"> 92
</span><span class="lnt"> 93
</span><span class="lnt"> 94
</span><span class="lnt"> 95
</span><span class="lnt"> 96
</span><span class="lnt"> 97
</span><span class="lnt"> 98
</span><span class="lnt"> 99
</span><span class="lnt">100
</span><span class="lnt">101
</span><span class="lnt">102
</span><span class="lnt">103
</span><span class="lnt">104
</span><span class="lnt">105
</span><span class="lnt">106
</span><span class="lnt">107
</span><span class="lnt">108
</span><span class="lnt">109
</span><span class="lnt">110
</span><span class="lnt">111
</span></code></pre></td>
<td class="lntd">
<pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">Dataset</span>
</span></span><span class="line"><span class="cl"><span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">LongTensor</span>
</span></span><span class="line"><span class="cl"><span class="kn">from</span> <span class="nn">transformers</span> <span class="kn">import</span> <span class="n">PreTrainedTokenizerFast</span>
</span></span><span class="line"><span class="cl"><span class="kn">import</span> <span class="nn">pyarrow.parquet</span> <span class="k">as</span> <span class="nn">pq</span>
</span></span><span class="line"><span class="cl"><span class="kn">from</span> <span class="nn">numpy</span> <span class="kn">import</span> <span class="n">array</span><span class="p">,</span> <span class="n">int64</span>
</span></span><span class="line"><span class="cl"><span class="kn">from</span> <span class="nn">numpy.random</span> <span class="kn">import</span> <span class="n">shuffle</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="k">class</span> <span class="nc">MyDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> 
</span></span><span class="line"><span class="cl">                <span class="n">parquet_file</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="n">tokenizer_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="n">keep_in_memory</span><span class="p">:</span> <span class="nb">bool</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="n">max_seq_len</span><span class="p">:</span> <span class="nb">int</span><span class="o">=</span><span class="mi">512</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="n">buffer_size</span><span class="p">:</span> <span class="nb">int</span><span class="o">=</span><span class="mi">40960</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">        <span class="s1">&#39;&#39;&#39;
</span></span></span><span class="line"><span class="cl"><span class="s1">        keep_in_memory: 是否将parquet文件转换为pandas.DataFrame格式存放到内存, 
</span></span></span><span class="line"><span class="cl"><span class="s1">            False将使用迭代生成器(迭代生成器不支持打乱数据)，减少大数据集内存占用
</span></span></span><span class="line"><span class="cl"><span class="s1">        &#39;&#39;&#39;</span>
</span></span><span class="line"><span class="cl">        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">keep_in_memory</span> <span class="o">=</span> <span class="n">keep_in_memory</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">max_seq_len</span> <span class="o">=</span> <span class="n">max_seq_len</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="c1"># 使用pyarrow.parquet读取，to_pandas、for遍历速度更快</span>
</span></span><span class="line"><span class="cl">        <span class="n">parquet_table</span> <span class="o">=</span> <span class="n">pq</span><span class="o">.</span><span class="n">read_table</span><span class="p">(</span><span class="n">parquet_file</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="c1"># 获取数据集长度</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">length</span> <span class="o">=</span> <span class="n">parquet_table</span><span class="o">.</span><span class="n">num_rows</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="c1"># 缓冲区大小不能超过数据长度</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">buffer_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">length</span> <span class="k">if</span> <span class="n">buffer_size</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">length</span> <span class="k">else</span> <span class="n">buffer_size</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="k">if</span> <span class="n">keep_in_memory</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">            <span class="c1"># 转化为pandas放到内存中</span>
</span></span><span class="line"><span class="cl">            <span class="bp">self</span><span class="o">.</span><span class="n">data</span> <span class="o">=</span> <span class="n">parquet_table</span><span class="o">.</span><span class="n">to_pandas</span><span class="p">()</span>  
</span></span><span class="line"><span class="cl">        <span class="k">else</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">            <span class="bp">self</span><span class="o">.</span><span class="n">data</span> <span class="o">=</span> <span class="n">parquet_table</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="c1"># 初始化tokenizer</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">tokenizer</span> <span class="o">=</span> <span class="n">PreTrainedTokenizerFast</span><span class="o">.</span><span class="n">from_pretrained</span><span class="p">(</span><span class="n">tokenizer_dir</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="c1"># 在这里初始化generator</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">sample_generator</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">item_generator</span><span class="p">()</span>
</span></span><span class="line"><span class="cl">    
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="nf">item_generator</span><span class="p">(</span><span class="bp">self</span><span class="p">,)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">        <span class="s1">&#39;&#39;&#39;
</span></span></span><span class="line"><span class="cl"><span class="s1">        一条数据的生成器，防止大数据集OOM
</span></span></span><span class="line"><span class="cl"><span class="s1">        &#39;&#39;&#39;</span>
</span></span><span class="line"><span class="cl">                
</span></span><span class="line"><span class="cl">        <span class="n">parquet_table</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">data</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="c1"># 生成器是死循环，不用退出，训练结束（epoch结束）会停止调用next()</span>
</span></span><span class="line"><span class="cl">        <span class="n">buffer_list</span> <span class="o">=</span> <span class="p">[]</span>
</span></span><span class="line"><span class="cl">        <span class="k">while</span> <span class="kc">True</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">            <span class="k">for</span> <span class="n">prompt</span><span class="p">,</span> <span class="n">response</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">parquet_table</span><span class="p">[</span><span class="s1">&#39;prompt&#39;</span><span class="p">],</span> <span class="n">parquet_table</span><span class="p">[</span><span class="s1">&#39;response&#39;</span><span class="p">]):</span>
</span></span><span class="line"><span class="cl">                
</span></span><span class="line"><span class="cl">                <span class="c1"># 缓存数据不够，添加数据</span>
</span></span><span class="line"><span class="cl">                <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">buffer_list</span><span class="p">)</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">buffer_size</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">                    <span class="n">buffer_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span> <span class="p">(</span><span class="n">prompt</span><span class="o">.</span><span class="n">as_py</span><span class="p">(),</span> <span class="n">response</span><span class="o">.</span><span class="n">as_py</span><span class="p">())</span> <span class="p">)</span>
</span></span><span class="line"><span class="cl">                    <span class="k">continue</span>
</span></span><span class="line"><span class="cl">                
</span></span><span class="line"><span class="cl">                <span class="c1"># 执行到这里，缓存区够了，打乱数据</span>
</span></span><span class="line"><span class="cl">                <span class="n">shuffle</span><span class="p">(</span><span class="n">buffer_list</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">                <span class="k">for</span> <span class="n">p</span><span class="p">,</span> <span class="n">r</span> <span class="ow">in</span> <span class="n">buffer_list</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">                    <span class="c1"># 在这里迭代</span>
</span></span><span class="line"><span class="cl">                    <span class="k">yield</span>  <span class="n">p</span><span class="p">,</span> <span class="n">r</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">                <span class="c1"># 迭代完成，清空缓存区</span>
</span></span><span class="line"><span class="cl">                <span class="n">buffer_list</span> <span class="o">=</span> <span class="p">[]</span>
</span></span><span class="line"><span class="cl">    
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="fm">__getitem__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">        <span class="s1">&#39;&#39;&#39;
</span></span></span><span class="line"><span class="cl"><span class="s1">        返回一条样本
</span></span></span><span class="line"><span class="cl"><span class="s1">        &#39;&#39;&#39;</span>
</span></span><span class="line"><span class="cl">        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">keep_in_memory</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">            <span class="n">data</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">data</span>
</span></span><span class="line"><span class="cl">            <span class="n">prompt</span><span class="p">,</span> <span class="n">response</span> <span class="o">=</span> <span class="n">data</span><span class="o">.</span><span class="n">iloc</span><span class="p">[</span><span class="n">index</span><span class="p">]</span><span class="o">.</span><span class="n">prompt</span><span class="p">,</span> <span class="n">data</span><span class="o">.</span><span class="n">iloc</span><span class="p">[</span><span class="n">index</span><span class="p">]</span><span class="o">.</span><span class="n">response</span>
</span></span><span class="line"><span class="cl">        <span class="k">else</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">            <span class="n">prompt</span><span class="p">,</span> <span class="n">response</span> <span class="o">=</span> <span class="nb">next</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">sample_generator</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="n">max_seq_len</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_seq_len</span> <span class="o">-</span> <span class="mi">5</span> <span class="c1"># len(&#39;[EOS]&#39;) = 5</span>
</span></span><span class="line"><span class="cl">        <span class="c1"># add an eos token note that end of resopnse, using in generate.</span>
</span></span><span class="line"><span class="cl">        <span class="k">return</span> <span class="sa">f</span><span class="s2">&#34;</span><span class="si">{</span><span class="n">prompt</span><span class="p">[</span><span class="mi">0</span><span class="p">:</span> <span class="n">max_seq_len</span><span class="p">]</span><span class="si">}</span><span class="s2">[EOS]&#34;</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&#34;</span><span class="si">{</span><span class="n">response</span><span class="p">[</span><span class="mi">0</span><span class="p">:</span> <span class="n">max_seq_len</span><span class="p">]</span><span class="si">}</span><span class="s2">[EOS]&#34;</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="nf">collate_fn</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="nb">list</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">        <span class="s1">&#39;&#39;&#39;
</span></span></span><span class="line"><span class="cl"><span class="s1">        合并一个批次数据返回
</span></span></span><span class="line"><span class="cl"><span class="s1">        &#39;&#39;&#39;</span>
</span></span><span class="line"><span class="cl">        <span class="n">tokenizer</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">tokenizer</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="n">prompt</span> <span class="o">=</span> <span class="n">tokenizer</span><span class="p">([</span><span class="n">item</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">data</span><span class="p">],</span> <span class="n">padding</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">return_token_type_ids</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">response</span> <span class="o">=</span> <span class="n">tokenizer</span><span class="p">([</span><span class="n">item</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">data</span><span class="p">],</span> <span class="n">padding</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">return_token_type_ids</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="n">input_ids</span> <span class="o">=</span> <span class="n">array</span><span class="p">(</span><span class="n">prompt</span><span class="o">.</span><span class="n">input_ids</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int64</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">input_mask</span> <span class="o">=</span> <span class="n">array</span><span class="p">(</span><span class="n">prompt</span><span class="o">.</span><span class="n">attention_mask</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int64</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">target_ids</span> <span class="o">=</span> <span class="n">array</span><span class="p">(</span><span class="n">response</span><span class="o">.</span><span class="n">input_ids</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int64</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="n">ret</span> <span class="o">=</span> <span class="p">{</span>
</span></span><span class="line"><span class="cl">            <span class="s1">&#39;input_ids&#39;</span><span class="p">:</span> <span class="n">LongTensor</span><span class="p">(</span><span class="n">input_ids</span><span class="p">),</span>
</span></span><span class="line"><span class="cl">            <span class="s1">&#39;input_mask&#39;</span><span class="p">:</span> <span class="n">LongTensor</span><span class="p">(</span><span class="n">input_mask</span><span class="p">),</span>
</span></span><span class="line"><span class="cl">            <span class="s1">&#39;target_ids&#39;</span><span class="p">:</span> <span class="n">LongTensor</span><span class="p">(</span><span class="n">target_ids</span><span class="p">),</span>
</span></span><span class="line"><span class="cl">        <span class="p">}</span>
</span></span><span class="line"><span class="cl">        <span class="k">return</span> <span class="n">ret</span>
</span></span><span class="line"><span class="cl">    
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="fm">__len__</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">length</span>
</span></span></code></pre></td></tr></table>
</div>
</div>
</section>


    <footer class="article-footer">
    

    </footer>


    
        <link 
                rel="stylesheet" 
                href="https://cdn.jsdelivr.net/npm/katex@0.13.13/dist/katex.min.css"integrity="sha384-RZU/ijkSsFbcmivfdRBQDtwuwVqK7GMOw6IMvKyeWL2K5UAlyp6WonmB8m7Jd0Hn"crossorigin="anonymous"
            ><script 
                src="https://cdn.jsdelivr.net/npm/katex@0.13.13/dist/katex.min.js"integrity="sha384-pK1WpvzWVBQiP0/GjnvRxV4mOb0oxFuyRxJlk6vVw146n3egcN5C925NCP7a7BY8"crossorigin="anonymous"
                defer
                >
            </script><script 
                src="https://cdn.jsdelivr.net/npm/katex@0.13.13/dist/contrib/auto-render.min.js"integrity="sha384-vZTG03m&#43;2yp6N6BNi5iM4rW4oIwk5DfcNdFfxkk9ZWpDriOkXX8voJBFrAO7MpVl"crossorigin="anonymous"
                defer
                >
            </script><script>
    window.addEventListener("DOMContentLoaded", () => {
        renderMathInElement(document.querySelector(`.article-content`), {
            delimiters: [
                { left: "$$", right: "$$", display: true },
                { left: "$", right: "$", display: false },
                { left: "\\(", right: "\\)", display: false },
                { left: "\\[", right: "\\]", display: true }
            ]
        });})
</script>
    
</article>

    

    

<aside class="related-contents--wrapper">
    <h2 class="section-title">相关文章</h2>
    <div class="related-contents">
        <div class="flex article-list--tile">
            
                
<article class="">
    <a href="/p/%E6%B7%B1%E5%BA%A6%E5%AD%A6%E4%B9%A0%E5%8F%8D%E5%90%91%E4%BC%A0%E6%92%AD%E6%8E%A8%E5%AF%BC/">
        
        

        <div class="article-details">
            <h2 class="article-title">深度学习反向传播推导</h2>
        </div>
    </a>
</article>
            
                
<article class="">
    <a href="/p/pytorch%E5%A4%A7%E6%95%B0%E6%8D%AE%E9%9B%86num_workers%E5%A4%A7%E4%BA%8E1%E5%86%85%E5%AD%98%E7%BC%93%E6%85%A2%E5%A2%9E%E6%B6%A8/">
        
        

        <div class="article-details">
            <h2 class="article-title">Pytorch大数据集num_workers大于1内存缓慢增涨</h2>
        </div>
    </a>
</article>
            
                
<article class="">
    <a href="/p/%E5%B0%8F%E6%A0%B7%E6%9C%AC%E5%AD%A6%E4%B9%A0/">
        
        

        <div class="article-details">
            <h2 class="article-title">小样本学习</h2>
        </div>
    </a>
</article>
            
                
<article class="">
    <a href="/p/%E5%AD%AA%E7%94%9F%E7%BD%91%E7%BB%9C%E5%AE%9E%E7%8E%B0%E5%A4%9A%E5%88%86%E7%B1%BBsiamese-network/">
        
        

        <div class="article-details">
            <h2 class="article-title">孪生网络实现多分类（Siamese Network）</h2>
        </div>
    </a>
</article>
            
                
<article class="">
    <a href="/p/%E4%BD%BF%E7%94%A8%E4%BA%8C%E5%88%86%E7%B1%BB%E5%AE%9E%E7%8E%B0%E5%A4%9A%E5%88%86%E7%B1%BB/">
        
        

        <div class="article-details">
            <h2 class="article-title">使用二分类实现多分类</h2>
        </div>
    </a>
</article>
            
        </div>
    </div>
</aside>

     
    
        
    

    <footer class="site-footer">
    <section class="copyright">
        &copy; 
        
            2018 - 
        
        2024 Charent的博客
    </section>
    
    <section class="powerby">
         <br />
        
    </section>
</footer>


    
<div class="pswp" tabindex="-1" role="dialog" aria-hidden="true">

    
    <div class="pswp__bg"></div>

    
    <div class="pswp__scroll-wrap">

        
        <div class="pswp__container">
            <div class="pswp__item"></div>
            <div class="pswp__item"></div>
            <div class="pswp__item"></div>
        </div>

        
        <div class="pswp__ui pswp__ui--hidden">

            <div class="pswp__top-bar">

                

                <div class="pswp__counter"></div>

                <button class="pswp__button pswp__button--close" title="Close (Esc)"></button>

                <button class="pswp__button pswp__button--share" title="Share"></button>

                <button class="pswp__button pswp__button--fs" title="Toggle fullscreen"></button>

                <button class="pswp__button pswp__button--zoom" title="Zoom in/out"></button>

                
                
                <div class="pswp__preloader">
                    <div class="pswp__preloader__icn">
                        <div class="pswp__preloader__cut">
                            <div class="pswp__preloader__donut"></div>
                        </div>
                    </div>
                </div>
            </div>

            <div class="pswp__share-modal pswp__share-modal--hidden pswp__single-tap">
                <div class="pswp__share-tooltip"></div>
            </div>

            <button class="pswp__button pswp__button--arrow--left" title="Previous (arrow left)">
            </button>

            <button class="pswp__button pswp__button--arrow--right" title="Next (arrow right)">
            </button>

            <div class="pswp__caption">
                <div class="pswp__caption__center"></div>
            </div>

        </div>

    </div>

</div><script 
                src="https://cdn.jsdelivr.net/npm/photoswipe@4.1.3/dist/photoswipe.min.js"integrity="sha256-ePwmChbbvXbsO02lbM3HoHbSHTHFAeChekF1xKJdleo="crossorigin="anonymous"
                defer
                >
            </script><script 
                src="https://cdn.jsdelivr.net/npm/photoswipe@4.1.3/dist/photoswipe-ui-default.min.js"integrity="sha256-UKkzOn/w1mBxRmLLGrSeyB4e1xbrp4xylgAWb3M42pU="crossorigin="anonymous"
                defer
                >
            </script><link 
                rel="stylesheet" 
                href="https://cdn.jsdelivr.net/npm/photoswipe@4.1.3/dist/default-skin/default-skin.css"integrity="sha256-c0uckgykQ9v5k&#43;IqViZOZKc47Jn7KQil4/MP3ySA3F8="crossorigin="anonymous"
            ><link 
                rel="stylesheet" 
                href="https://cdn.jsdelivr.net/npm/photoswipe@4.1.3/dist/photoswipe.css"integrity="sha256-SBLU4vv6CA6lHsZ1XyTdhyjJxCjPif/TRkjnsyGAGnE="crossorigin="anonymous"
            >

            </main>
    

        </div>
        <script 
                src="https://cdn.jsdelivr.net/npm/node-vibrant@3.1.5/dist/vibrant.min.js"integrity="sha256-5NovOZc4iwiAWTYIFiIM7DxKUXKWvpVEuMEPLzcm5/g="crossorigin="anonymous"
                
                >
            </script><script type="text/javascript" src="/ts/main.js" defer></script>
<script>
    (function () {
        const customFont = document.createElement('link');
        customFont.href = "https://fonts.googleapis.com/css2?family=Lato:wght@300;400;700&display=swap";

        customFont.type = "text/css";
        customFont.rel = "stylesheet";

        document.head.appendChild(customFont);
    }());
</script>

    </body>
</html>
